import time
from ._window import Window
from ._window_data import PartialWindowData
from torch.utils.data.dataloader import DataLoader
import torch


class TimeWindow(Window):
    def __init__(self, data_loader: DataLoader, window_size: int, threshold: int):
        super().__init__()
        self.data_loader = iter(data_loader)
        self.window_size = window_size
        self.threshold = threshold
        self.last_time = None

    @staticmethod
    def get_secs():
        return int(time.time())

    def __inner_poll(self):
        loader_consumed = False
        cached_X = []
        cached_y = []
        cached_size = 0
        trigger_window = False
        for batch_data in self.data_loader:
            if batch_data is None:
                loader_consumed = True
                break
            X, y = batch_data

            cached_X.append(X)
            cached_y.append(y)
            cached_size += len(X)
            self.offset += len(X)

            if self.last_time + self.window_size >= TimeWindow.get_secs():
                trigger_window = True
                break

            if cached_size >= self.threshold:
                break

        if loader_consumed:
            return None

        if trigger_window:
            self.last_time = TimeWindow.get_secs()

        X = torch.cat(cached_X, dim=0)
        y = torch.cat(cached_y, dim=0)

        window_data = PartialWindowData(X, y, trigger_window)
        return window_data

    def poll(self) -> PartialWindowData:
        if self.last_time is None:
            self.last_time = TimeWindow.get_secs()
        return self.__inner_poll()
